package MLDA_Process;

import gnu.trove.map.hash.TIntIntHashMap;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import Commons.IOUtil;

public class PAM4ST_V3 implements Runnable {

    public int P = 1000;
    public int K = 100;

    public double alpha1 = 0.01;
    public double alpha2 = 0.1;
    public double beta = 0.01;

    public double vBeta;
    public double pAlpha1;
    public double kAlpha2;

    public int niters = 5000;

    public int M;
    public int V;

    public int[] nkwSum;
    public int[] npkSum;
    public int[] ndpSum;

    public TIntIntHashMap[] nwk;
    public TIntIntHashMap[] npk;
    public TIntIntHashMap[] ndp;

    // sample z1
    public double e_sum1;
    public double f_sum1;
    public double g_sum1;

    public double[] c1; // c1[p][k] = alpha2 / (npkSum[p] + pAlpha1)

    // sample z2
    public double e_sum2;
    public double f_sum2;
    public double g_sum2;

    public double[] c2; // c2[k] = alpha1 / (nkwSum[k] + vBeta)

    public int[][] z1_assigns;
    public int[][] z2_assigns;

    public List<List<Integer>> docs = new ArrayList<List<Integer>>();
    public HashMap<String, Integer> w2i = new HashMap<String, Integer>();
    public HashMap<Integer, String> i2w = new HashMap<Integer, String>();

    public int save_step = 1000;
    public String inputPath = "";
    public String outputPath = "";

    public PAM4ST_V3(int P, int K, int niters, int save_step, String inputPath, String outputPath) {
        this.P = P;
        this.save_step = save_step;
        this.K = K;
        this.niters = niters;
        this.inputPath = inputPath;
        this.outputPath = outputPath;
    }

    public void loadTxts(String txtPath) {
        BufferedReader reader = IOUtil.getReader(txtPath, "UTF-8");

        String line;
        try {
            line = reader.readLine();
            while (line != null) {
                List<Integer> doc = new ArrayList<Integer>();

                String[] tokens = line.trim().split("\\s+");
                for (String token : tokens) {
                    if (!w2i.containsKey(token)) {
                        w2i.put(token, w2i.size());
                        i2w.put(w2i.get(token), token);
                    }
                    doc.add(w2i.get(token));
                }
                docs.add(doc);
                line = reader.readLine();
            }
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }

        M = docs.size();
        V = w2i.size();
        return;
    }

    public void init_model() {
        vBeta = V * beta;
        pAlpha1 = P * alpha1;
        kAlpha2 = K * alpha2;

        nkwSum = new int[K];
        npkSum = new int[P];
        ndpSum = new int[M];

        nwk = new TIntIntHashMap[V];
        for (int v = 0; v != V; v++)
            nwk[v] = new TIntIntHashMap();
        npk = new TIntIntHashMap[P];
        for (int p = 0; p != P; p++)
            npk[p] = new TIntIntHashMap();
        ndp = new TIntIntHashMap[M];
        for (int m = 0; m != M; m++)
            ndp[m] = new TIntIntHashMap();

        z1_assigns = new int[M][];
        z2_assigns = new int[M][];

        for (int m = 0; m != M; m++) {
            int N = docs.get(m).size();
            z1_assigns[m] = new int[N];
            z2_assigns[m] = new int[N];

            for (int n = 0; n != N; n++) {
                int w = docs.get(m).get(n);
                int z1 = (int) Math.floor(Math.random() * P);
                int z2 = (int) Math.floor(Math.random() * K);

                nkwSum[z2]++;
                npkSum[z1]++;
                ndpSum[m]++;

                nwk[w].adjustOrPutValue(z2, 1, 1);
                npk[z1].adjustOrPutValue(z2, 1, 1);
                ndp[m].adjustOrPutValue(z1, 1, 1);

                z1_assigns[m][n] = z1;
                z2_assigns[m][n] = z2;
            }
        }

        c1 = new double[P];
        for (int p = 0; p != P; p++) {
            c1[p] = alpha2 / (npkSum[p] + kAlpha2);
        }

        g_sum1 = 0.0;
        for (int p = 0; p != P; p++) {
            g_sum1 += alpha1 * c1[p];
        }

        c2 = new double[K];
        for (int k = 0; k != K; k++) {
            c2[k] = alpha2 / (nkwSum[k] + vBeta);
        }

        g_sum2 = 0.0;
        for (int k = 0; k != K; k++) {
            g_sum2 += alpha2 * c2[k];
        }
    }

    public void sample_doc_z1(int m) {

        for (int n = 0; n != docs.get(m).size(); n++) {
            int z1 = z1_assigns[m][n];
            int z2 = z2_assigns[m][n];

            g_sum1 -= alpha1 * c1[z1];

            if (npk[z1].get(z2) == 1) {
                npk[z1].remove(z2);
            } else {
                npk[z1].adjustValue(z2, -1);
            }
            npkSum[z1]--;

            if (ndp[m].get(z1) == 1) {
                ndp[m].remove(z1);
            } else {
                ndp[m].adjustValue(z1, -1);
            }
            ndpSum[m]--;

            c1[z1] = alpha2 / (npkSum[z1] + kAlpha2);
            g_sum1 += alpha1 * c1[z1];

            e_sum1 = 0.0;
            for (int p : ndp[m].keys()) {
                if (npk[p].contains(z2))
                    e_sum1 += ndp[m].get(p) * (npk[p].get(z2) / (npkSum[p] + kAlpha2) + c1[p]);
                else
                    e_sum1 += ndp[m].get(p) * c1[p];
            }

            f_sum1 = 0.0;
            for (int p = 0; p != P; p++) {
                if (npk[p].contains(z2)) {
                    f_sum1 += npk[p].get(z2) * alpha1 / (npkSum[p] + kAlpha2);
                }
            }

            double total_mass = e_sum1 + f_sum1 + g_sum1;
            double u = Math.random() * total_mass;
            double r = 0.0, p_sum = 0.0;

            if (u < e_sum1) {
                r = Math.random() * e_sum1;
                for (int p : ndp[m].keys()) {
                    if (npk[p].contains(z2))
                        p_sum += ndp[m].get(p) * (npk[p].get(z2) / (npkSum[p] + kAlpha2) + c1[p]);
                    else
                        p_sum += ndp[m].get(p) * c1[p];

                    if (r < p_sum) {
                        z1 = p;
                        break;
                    }
                }

            } else {
                if (u < e_sum1 + f_sum1) {
                    r = Math.random() * f_sum1;
                    for (int p = 0; p != P; p++) {
                        if (npk[p].contains(z2)) {
                            p_sum += npk[p].get(z2) * alpha1 / (npkSum[p] + kAlpha2);
                            if (r < p_sum) {
                                z1 = p;
                                break;
                            }
                        }
                    }
                } else {
                    r = Math.random() * g_sum1;
                    for (int p = 0; p != P; p++) {
                        p_sum += alpha1 * c1[p];
                        if (r < p_sum) {
                            z1 = p;
                            break;
                        }
                    }
                }
            }

            //

            g_sum1 -= alpha1 * c1[z1];

            npk[z1].adjustOrPutValue(z2, 1, 1);
            npkSum[z1]++;

            ndp[m].adjustOrPutValue(z1, 1, 1);
            ndpSum[m]++;

            c1[z1] = alpha2 / (npkSum[z1] + kAlpha2);
            g_sum1 += alpha1 * c1[z1];

            z1_assigns[m][n] = z1;
        }
    }

    public void sample_doc_z2(int m) {

        for (int n = 0; n != docs.get(m).size(); n++) {
            int z1 = z1_assigns[m][n];
            int z2 = z2_assigns[m][n];
            int w = docs.get(m).get(n);

            g_sum2 -= beta * c2[z2];

            if (nwk[w].get(z2) == 1) {
                nwk[w].remove(z2);
            } else {
                nwk[w].adjustValue(z2, -1);
            }
            nkwSum[z2]--;

            if (npk[z1].get(z2) == 1) {
                npk[z1].remove(z2);
            } else {
                npk[z1].adjustValue(z2, -1);
            }
            npkSum[z1]--;

            c2[z2] = alpha2 / (nkwSum[z2] + vBeta);
            g_sum2 += beta * c2[z2];

            e_sum2 = 0.0;
            for (int k : nwk[w].keys()) {
                if (npk[z1].contains(k)) {
                    e_sum2 += nwk[w].get(k) * (npk[z1].get(k) / (nkwSum[k] + vBeta) + c2[k]);
                } else {
                    e_sum2 += nwk[w].get(k) * c2[k];
                }
            }

            f_sum2 = 0.0;
            for (int k = 0; k != K; k++) {
                if (npk[z1].contains(k)) {
                    f_sum2 += npk[z1].get(k) * beta / (nkwSum[k] + vBeta);
                }
            }

            double total_mass = e_sum2 + f_sum2 + g_sum2;
            double u = Math.random() * total_mass;
            double r = 0.0, p_sum = 0.0;

            if (u < e_sum2) {
                r = Math.random() * e_sum2;
                for (int k : nwk[w].keys()) {
                    if (npk[z1].contains(k)) {
                        p_sum += nwk[w].get(k) * (npk[z1].get(k) / (nkwSum[k] + vBeta) + c2[k]);
                    } else {
                        p_sum += nwk[w].get(k) * c2[k];
                    }
                    if (r < p_sum) {
                        z2 = k;
                        break;
                    }
                }
            } else {
                if (u < e_sum2 + f_sum2) {
                    r = Math.random() * f_sum2;
                    for (int k = 0; k != K; k++) {
                        if (npk[z1].contains(k)) {
                            p_sum += npk[z1].get(k) * beta / (nkwSum[k] + vBeta);
                            if (r < p_sum) {
                                z2 = k;
                                break;
                            }
                        }
                    }
                } else {
                    r = Math.random() * g_sum2;
                    for (int k = 0; k != K; k++) {
                        p_sum += alpha2 * c2[k];
                        if (r < p_sum) {
                            z2 = k;
                            break;
                        }
                    }
                }
            }

            //

            g_sum2 -= beta * c2[z2];

            nwk[w].adjustOrPutValue(z2, 1, 1);
            nkwSum[z2]++;

            npk[z1].adjustOrPutValue(z2, 1, 1);
            npkSum[z1]++;

            c2[z2] = alpha2 / (nkwSum[z2] + vBeta);
            g_sum2 += beta * c2[z2];

            z2_assigns[m][n] = z2;
        }

    }

    public void compact_hashtables() {
        System.out.println("Compacting tables ...");
        for (int m = 0; m != M; m++) {
            ndp[m].compact();
        }

        for (int v = 0; v != V; v++) {
            nwk[v].compact();
        }

        for (int p = 0; p != P; p++) {
            npk[p].compact();
        }
    }

    public void estimate() {
        for (int iter = 0; iter != niters; iter++) {
            System.out.println("PAM4ST_V3 Iteration " + iter + " ...");
            if (iter % this.save_step == 0 && iter != 0) {
                try {
                    this.printTopics(outputPath + "/topTopics_" + iter + ".txt", 20);
                    this.saveWordmap(outputPath + "/wordmap_" + iter + ".txt");
                    this.savePhi(outputPath + "/phi_" + iter + ".txt");
                    this.saveZAssigns(outputPath + "/zassign_" + iter + ".txt");
                    this.saveTheta(outputPath + "/theta_" + iter + ".txt");
                } catch (IOException e) {
                    // TODO Auto-generated catch block
                    e.printStackTrace();
                }
            }

            long begin = System.currentTimeMillis();
            for (int m = 0; m != M; m++) {
                sample_doc_z1(m);
            }
            for (int m = 0; m != M; m++) {
                sample_doc_z2(m);
            }
            System.out.println((System.currentTimeMillis() - begin) / 1000.0);
            if ((iter + 1) % 100 == 0)
                compact_hashtables();
        }
    }

    int nkw[][];
    int ndk[][];
    int ndkSum[];

    public void convert_map_to_arrays_phi() {
        nkw = new int[K][V];
        for (int w = 0; w != V; w++) {
            TIntIntHashMap nwkTMap = nwk[w];
            for (int k : nwkTMap.keys()) {
                nkw[k][w] = nwkTMap.get(k);
            }
        }
    }

    public void convert_zassigns_to_arrays_theta() {
        ndk = new int[M][K];
        ndkSum = new int[M];

        for (int m = 0; m != M; m++) {
            for (int n = 0; n != docs.get(m).size(); n++) {
                ndk[m][z2_assigns[m][n]]++;
                ndkSum[m]++;
            }
        }
    }

    public double[][] computeTheta() {
        convert_zassigns_to_arrays_theta();
        double[][] theta = new double[M][K];
        for (int m = 0; m != M; m++) {
            for (int k = 0; k != K; k++) {
                theta[m][k] = (ndk[m][k] + alpha2) / (ndkSum[m] + K * alpha2);
            }
        }
        return theta;
    }

    public double[][] computePhi() {
        convert_map_to_arrays_phi();
        double[][] phi = new double[K][V];
        for (int k = 0; k != K; k++) {
            for (int v = 0; v != V; v++) {
                phi[k][v] = (nkw[k][v] + beta) / (nkwSum[k] + V * beta);
            }
        }
        return phi;
    }

    public ArrayList<List<Entry<String, Double>>> sortedTopicWords(
            double[][] phi, int T) {
        ArrayList<List<Entry<String, Double>>> res = new ArrayList<List<Entry<String, Double>>>();
        for (int k = 0; k != T; k++) {
            HashMap<String, Double> term2weight = new HashMap<String, Double>();
            for (String term : w2i.keySet())
                term2weight.put(term, phi[k][w2i.get(term)]);

            List<Entry<String, Double>> pairs = new ArrayList<Entry<String, Double>>(
                    term2weight.entrySet());
            Collections.sort(pairs, new Comparator<Entry<String, Double>>() {
                public int compare(Entry<String, Double> o1,
                                   Entry<String, Double> o2) {
                    return (o2.getValue().compareTo(o1.getValue()));
                }
            });
            res.add(pairs);
        }
        return res;
    }

    public void printTopics(String path, int top_n) throws IOException {
        BufferedWriter writer = IOUtil.getWriter(path);
        double[][] phi = computePhi();
        ArrayList<List<Entry<String, Double>>> pairsList = this
                .sortedTopicWords(phi, K);
        for (int k = 0; k != K; k++) {
            writer.write("Topic " + k + ":\n");
            for (int i = 0; i != top_n; i++) {
                writer.write(pairsList.get(k).get(i).getKey() + " "
                        + pairsList.get(k).get(i).getValue() + "\n");
            }
        }
        writer.close();
    }

    public void savePhi(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] phi = computePhi();
        try {
            for (int k = 0; k != K; k++) {
                for (int v = 0; v != V; v++)
                    writer.append(phi[k][v] + " ");
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveTheta(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        double[][] theta = computeTheta();
        try {
            for (int m = 0; m != M; m++) {
                for (int k = 0; k != K; k++)
                    writer.append(theta[m][k] + " ");
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveZAssigns(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (int m = 0; m != M; m++) {
                for (int n = 0; n != docs.get(m).size(); n++)
                    writer.append(z1_assigns[m][n] + ":" + z2_assigns[m][n] + ":" + docs.get(m).get(n) + " ");
                writer.append("\n");
            }
            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void saveWordmap(String path) {
        BufferedWriter writer = IOUtil.getWriter(path, "utf-8");

        try {
            for (String word : w2i.keySet())
                writer.append(word + "\t" + w2i.get(word) + "\n");

            writer.flush();
            writer.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return;
    }

    public void printModel() {
        System.out.println("save_step:" + this.save_step + "\tniters:" + this.niters + "\tbeta:" + this.beta + "\talpha1:" + this.alpha1 + "\talpha2:" + this.alpha2 + "\tk1:" + this.P + "\tk2:" + this.K + "\tinput:" + this.inputPath + "\toutput:" + this.outputPath);
    }

    public void run() {
        printModel();
        this.loadTxts(inputPath);//
        this.init_model();
        this.estimate();
        try {
            this.printTopics(outputPath + "/topTopics.txt", 20);
            this.saveWordmap(outputPath + "/wordmap.txt");
            this.savePhi(outputPath + "/phi.txt");
            this.saveZAssigns(outputPath + "/zassign.txt");
            this.saveTheta(outputPath + "/theta.txt");
        } catch (IOException e) {
            // TODO Auto-generated catch block
            e.printStackTrace();
        }

    }
//	public static void main(String args[]) {
//		PAM4ST_V3 pam4st = new PAM4ST_V3();
//		pam4st.loadTxts("/Users/zuoyuan/Desktop/Experiment/20130601-origin");//
//		pam4st.init_model();
//		pam4st.estimate();
//		pam4st.printTopics(10);
//	}

    public static void doPAM4ST_V3(int P, int K, int iter, int save_step, int threadNum, String path) {
        ExecutorService service = Executors.newCachedThreadPool();
        File trainFile = new File(path);
        String parent_path = trainFile.getParentFile().getAbsolutePath();
        for (int i = 0; i < threadNum; i++) {
            (new File(parent_path + "/PAM4ST_V3/" + i)).mkdirs();
            service.submit(new PAM4ST_V3(P, K, iter, save_step, path, parent_path + "/PAM4ST_V3/" + i + "/"));

        }
        service.shutdown();
    }

}
